import shap
# import sys
from prettytable import PrettyTable
from sklearn.ensemble import RandomForestRegressor as RFR
from sklearn.model_selection import train_test_split as TTS
# from sko.AFSA import AFSA
from streamlit_extras.colored_header import colored_header
from streamlit_shap import st_shap

from business.algorithm.utils import *


def run():
    colored_header(label="可解释机器学习", description=" ", color_name="violet-90")

    file = st.file_uploader("Upload `.csv`file", type=['csv'], label_visibility="collapsed")
    if file is None:
        table = PrettyTable(['file name', 'class', 'description'])
        table.add_row(['file_1', 'dataset', 'data file'])
        st.write(table)
    if file is not None:
        df = pd.read_csv(file)
        check_string_NaN(df)
        colored_header(label="数据信息", description=" ", color_name="violet-70")
        nrow = st.slider("rows", 1, len(df), 5)
        df_nrow = df.head(nrow)
        st.write(df_nrow)

        colored_header(label="特征&目标", description=" ", color_name="violet-70")

        target_num = st.number_input('目标数量', min_value=1, max_value=10, value=1)

        col_feature, col_target = st.columns(2)

        features = df.iloc[:, :-target_num]
        # targets
        targets = df.iloc[:, -target_num:]
        with col_feature:
            st.write(features.head())
        with col_target:
            st.write(targets.head())

        colored_header(label="Shapley value", description=" ", color_name="violet-70")

        fs = FeatureSelector(features, targets)

        target_selected_option = st.selectbox('choose target', list(fs.targets))
        fs.targets = fs.targets[target_selected_option]

        reg = RFR()
        X_train, X_test, y_train, y_test = TTS(fs.features, fs.targets, random_state=0)
        test_size = st.slider('test size', 0.1, 0.5, 0.2)
        random_state = st.checkbox('random state 42', True)
        if random_state:
            random_state = 42
        else:
            random_state = None

        fs.Xtrain, fs.Xtest, fs.Ytrain, fs.Ytest = TTS(fs.features, fs.targets, test_size=test_size,
                                                       random_state=random_state)
        reg.fit(fs.Xtrain, fs.Ytrain)

        explainer = shap.TreeExplainer(reg)

        shap_values = explainer(fs.features)

        colored_header(label="SHAP Feature Importance", description=" ", color_name="violet-30")
        nfeatures = st.slider("features", 2, fs.features.shape[1], fs.features.shape[1])
        st_shap(shap.plots.bar(shap_values, max_display=nfeatures))

        colored_header(label="SHAP Feature Cluster", description=" ", color_name="violet-30")
        clustering = shap.utils.hclust(fs.features, fs.targets)
        clustering_cutoff = st.slider('clustering cutoff', 0.0, 1.0, 0.5)
        nfeatures = st.slider("features", 2, fs.features.shape[1], fs.features.shape[1], key=2)
        st_shap(shap.plots.bar(shap_values, clustering=clustering, clustering_cutoff=clustering_cutoff,
                               max_display=nfeatures))

        colored_header(label="SHAP Beeswarm", description=" ", color_name="violet-30")
        rank_option = st.selectbox('rank option', ['max', 'mean'])
        max_dispaly = st.slider('max display', 2, fs.features.shape[1], fs.features.shape[1])
        if rank_option == 'max':
            st_shap(shap.plots.beeswarm(shap_values, order=shap_values.abs.max(0), max_display=max_dispaly))
        else:
            st_shap(shap.plots.beeswarm(shap_values, order=shap_values.abs.mean(0), max_display=max_dispaly))

        colored_header(label="SHAP Dependence", description=" ", color_name="violet-30")

        shap_values = explainer.shap_values(fs.features)
        list_features = fs.features.columns.tolist()
        feature = st.selectbox('特征', list_features)
        interact_feature = st.selectbox('interact feature', list_features)
        st_shap(shap.dependence_plot(feature, shap_values, fs.features, display_features=fs.features,
                                     interaction_index=interact_feature))